from sklearn.metrics import roc_auc_score
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from DenseGraphMatching import MultiLevelGraphMatchNetwork
from MGMN_dataset import MGMN_dataset
from cfg_config import malSim_args
from sklearn.metrics import auc, roc_curve
from datetime import datetime
import torch
import os
import numpy as np
from dataPre import dataPre
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")


def write_log_file(file_name_path, log_str, print_flag=True):
    if print_flag:
        print(log_str)
    if log_str is None:
        log_str = 'None'
    if os.path.isfile(file_name_path):
        with open(file_name_path, 'a+') as log_file:
            log_file.write(log_str + '\n')
    else:
        with open(file_name_path, 'w+') as log_file:
            log_file.write(log_str + '\n')


class MGMNTrainer(object):
    def __init__(self, model, args):
        super(MGMNTrainer, self).__init__()
        self.max_epoch = args.epochs
        self.batch_size = args.batch_size

        node_init_dims = 8

        log_file = r"/home/cyw/projects/function_sim_project/all_logs/MGMN_train_logs.txt"
        best_model_file = r"/home/cyw/projects/function_sim_project/all_data/models/MGMN_model.pth"
        self.lr = args.lr
        self.log_file = log_file
        self.best_model_path = best_model_file
        self.model = model

        write_log_file(self.log_file, "-----begin train-----", print_flag=True)
        write_log_file(self.log_file, str(self.model))

        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        write_log_file(self.log_file, "已加载数据集 @ {}".format(datetime.now()))

    def fit(self, dataloader, validloader, values):
        lossdata = []
        best_auc = -float("inf")
        print("MGMN train starting ~")
        for i in tqdm(range(1, self.max_epoch + 1)):
            self.model.train()
            cum_loss, num = 0.0, 0
            for j, (value, target1, target2) in enumerate(dataloader):
                torch.cuda.empty_cache()
                self.optimizer.zero_grad()

                adj_x, att_x = target1
                adj_y, att_y = target2

                adj_x = adj_x.cuda()
                adj_y = adj_y.cuda()
                att_x = att_x.cuda()
                att_y = att_y.cuda()
                value = value.cuda()

                batch_output = self.model(batch_x_p=att_x, batch_x_h=att_y,
                                          batch_adj_p=adj_x, batch_adj_h=adj_y)
                mse_loss = torch.nn.functional.mse_loss(batch_output, value)
                mse_loss.backward()
                self.optimizer.step()
                cum_loss += mse_loss
            temp_auc = self.test_model_auc(validloader, values)

            lossdata.append(cum_loss)
            write_log_file(self.log_file, "EPOCH {0}/{1}:\tMSE loss = {2}\t auc={3} @ {4}".
                           format(i, self.max_epoch, cum_loss, temp_auc, datetime.now()))
            if temp_auc > best_auc:
                best_auc = temp_auc
                #   参数不想调整了，就用论文里面的吧，训练到最优就行
                torch.save(
                    self.model, r'/home/cyw/projects/function_sim_project/all_data/models/MGMN_model.pth')
                write_log_file(self.log_file, "模型已保存", print_flag=True)
        torch.save(self.model, self.best_model_path)
        print("MGMN train finished!!")

    def get_model_dataloader(self, datas, name, is_shuffle=True):
        my_dataset = MGMN_dataset(datas, name)
        dataloader = DataLoader(my_dataset, batch_size=2, shuffle=is_shuffle, num_workers=32,
                                collate_fn=my_dataset.adjust_samples_to_same_dimension)
        return dataloader

    def test_model_auc(self, dataloader, values):
        ans = self.get_model_score(dataloader)
        auc_score = roc_auc_score(values, ans)
        return auc_score

    def get_model_score(self, dataloader):
        self.model.eval()
        with torch.no_grad():
            res = []
            for j, (value, target1, target2) in enumerate(dataloader):
                torch.cuda.empty_cache()
                self.optimizer.zero_grad()

                adj_x, att_x = target1
                adj_y, att_y = target2

                adj_x = adj_x.cuda()
                adj_y = adj_y.cuda()
                att_x = att_x.cuda()
                att_y = att_y.cuda()
                value = value.cuda()

                score1 = self.model(batch_x_p=att_x, batch_x_h=att_y,
                                    batch_adj_p=adj_x, batch_adj_h=adj_y)
                res.append(score1)
            ans = torch.cat((res))
        return ans.tolist()


if __name__ == '__main__':

    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:2000"

    nodeInitDim = 8

    datapre = dataPre()
    data_pair_infs = datapre.get_pair_infs()

    dataInf = data_pair_infs["train"]
    dataInf = [row[:3] for row in dataInf]  # 不要家族信息
    valid_dataInf = data_pair_infs["valid"]
    validvalues = []
    valid_dataInf = [row[:3] for row in valid_dataInf]  # 不要家族信息
    for i in range(len(data_pair_infs["valid"])):
        validvalues.append(data_pair_infs["valid"][i][2])

    # model = MultiLevelGraphMatchNetwork(node_init_dims=nodeInitDim, arguments=malSim_args).cuda()
    model = torch.load(
        r'/home/cyw/projects/function_sim_project/all_data/models/MGMN_model.pth')
    cfg_trainer = MGMNTrainer(model, args=malSim_args)
    dataloader = cfg_trainer.get_model_dataloader(dataInf, "train")
    valid_dataloader = cfg_trainer.get_model_dataloader(
        valid_dataInf, "valid", is_shuffle=False)
    ret_best_val_auc = cfg_trainer.fit(
        dataloader, valid_dataloader, validvalues)
